import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from model import LSTMModel
from data_process import prepare_stock_data


def predict_stock_prices(model, data_loader, scaler, device):
    """使用训练好的模型进行预测"""
    model.eval()
    predictions = []
    actuals = []

    with torch.no_grad():
        for X_batch, y_batch in data_loader:
            X_batch = X_batch.to(device)
            y_pred = model(X_batch)
            predictions.extend(y_pred.cpu().numpy())
            actuals.extend(y_batch.numpy())

    # 反向转换预测值到原始价格范围
    predictions = scaler.inverse_transform(
        np.array(predictions).reshape(-1, 1))
    actuals = scaler.inverse_transform(np.array(actuals).reshape(-1, 1))

    return predictions, actuals


def plot_predictions(predictions, actuals, title):
    """绘制预测结果"""
    plt.figure(figsize=(12, 6))
    plt.plot(actuals, label='Actual Prices')
    plt.plot(predictions, label='Predicted Prices')
    plt.title(title)
    plt.xlabel('Time')
    plt.ylabel('Stock Price')
    plt.legend()
    plt.show()


def main():
    # 加载模型
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = LSTMModel().to(device)
    model.load_state_dict(torch.load('stock_lstm_model.pth'))

    # 准备测试数据
    _, test_dataset, scaler = prepare_stock_data('stock_data.csv')
    test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)

    # 进行预测
    predictions, actuals = predict_stock_prices(
        model, test_loader, scaler, device)

    # 绘制结果
    plot_predictions(predictions, actuals, 'Stock Price Prediction')


if __name__ == '__main__':
    main()
